package com.nanohadoop.yarn;

import com.nanohadoop.mapreduce.JobConfiguration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

/**
 * YARN客户端，用于与YARN集群交互
 */
public class YARNClient {
    private static final Logger log = LoggerFactory.getLogger(YARNClient.class);

    private final ResourceManager resourceManager;
    private final Map<String, ApplicationMaster> applications;

    /**
     * 构造函数
     * @param resourceManager ResourceManager实例
     */
    public YARNClient(ResourceManager resourceManager) {
        this.resourceManager = resourceManager;
        this.applications = new HashMap<>();
    }

    /**
     * 提交MapReduce作业到YARN
     */
    public String submitApplication(JobConfiguration jobConf) throws IOException {
        // 确保ResourceManager已启动
        try {
            resourceManager.start();
        } catch (IOException e) {
            // 如果已启动，忽略异常
            if (!e.getMessage().contains("Address already in use")) {
                throw e;
            }
        }

        // 提交应用到ResourceManager
        ApplicationMaster appMaster = resourceManager.submitApplication(jobConf);
        String appId = appMaster.getApplicationId();

        // 存储ApplicationMaster引用
        applications.put(appId, appMaster);

        return appId;
    }

    /**
     * 获取ApplicationMaster
     */
    public ApplicationMaster getApplicationMaster(String appId) {
        return applications.get(appId);
    }

    /**
     * 停止应用
     */
    public boolean stopApplication(String appId) {
        ApplicationMaster appMaster = applications.get(appId);
        if (appMaster != null) {
            appMaster.setStatus(ApplicationMaster.ApplicationStatus.KILLED);
            log.debug("Application " + appId + " killed");
            return true;
        }
        return false;
    }
}